# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Face Recognition loss."""
import numpy as np
import mindspore
import mindspore.ops as ops
import math 
from mindspore import Tensor, Parameter
from mindspore.nn import Cell
from mindspore.nn.loss.loss import _Loss
from mindspore.ops import operations as P
from mindspore.ops import functional as F
from mindspore.common import dtype as mstype
from mindspore.ops import composite as C
import mindspore.common.dtype as mstype


__all__ = ['get_loss']

eps = 1e-24
def myphi(x,m):
    x = x * m
    return 1-x**2/math.factorial(2)+x**4/math.factorial(4)-x**6/math.factorial(6) + \
            x**8/math.factorial(8) - x**9/math.factorial(9)      
class AngleLinear(nn.Cell):
    def __init__(self, in_features, out_features, m = 4, phiflag=True):
        super(AngleLinear, self).__init__()
        mul = ops.Mul()
        
        self.in_features = in_features                                  
        self.out_features = out_features       
        self.norm = nn.Norm                   
        self.weight = Parameter(Tensor(in_features,out_features))      
        self.weight.data.uniform_(-1, 1).norm_(2,1,1e-5).mul_(1e5)        
       
        self.phiflag = phiflag
        self.m = m
        self.mlambda = [
            lambda x: x**0,                                                
            lambda x: x**1,
            lambda x: 2*x**2-1,
            lambda x: 4*x**3-3*x,
            lambda x: 8*x**4-8*x**2+1,
            lambda x: 16*x**5-20*x**3+5*x                     
        ]

    def forward(self, input):
        x = input   
        w = self.weight 
        floor = ops.Floor()
        acos = ops.ACos()
        min_value = Tensor(-1, mstype.float32)
        max_value = Tensor(1, mstype.float32)
        reshape = ops.Reshape()
        ww = w.norm(2,1,1e-5).mul(1e5)           
        xlen = x.pow(2).sum(1).pow(0.5)   
        wlen = ww.pow(2).sum(0).pow(0.5)      
        
        cos_theta = matmul(x, ww)               
        cos_theta = cos_theta / reshape(xlen, (-1, 1)) / reshape(xlen, (1, -1))   
        cos_theta = C.clip_by_value(cos_theta, min_value, max_value)           

        if self.phiflag:         
            cos_m_theta = self.mlambda[self.m](cos_theta)         
            theta = acos(cos_theta.data)               
            k = floor((self.m*theta/3.14159265))           
            n_one = k*0.0 - 1                                 
            phi_theta = (n_one**k) * cos_m_theta - 2*k      
        else:
            theta = acos(cos_theta)                   
            phi_theta = myphi(theta,self.m)
            phi_theta = C.clip_by_value(phi_theta, -1*self.m, 1)      

        cos_theta = cos_theta * reshape(xlen, (-1, 1)) 
        phi_theta = phi_theta * reshape(xlen, (-1, 1))
        output = (cos_theta,phi_theta)
        return output     


class AngleLoss(nn.Cell):
    def __init__(self, gamma=0):
        super(AngleLoss, self).__init__()
        self.gamma   = gamma
        self.it = 0
        self.LambdaMin = 5.0
        self.LambdaMax = 1500.0        
        self.lamb = 1500.0

    def forward(self, input, target):
        self.it += 1
        log_softmax = nn.LogSoftmax()
        cos_theta,phi_theta = input
        target =  reshape(target, (-1, 1))                     
      
        index = cos_theta.data * 0.0 
        op = ops.ScatterNd()
        indices = reshape(target.data, (-1, 1))
        updates = 1
        shape = index                          
        index = op(indices, updates, shape)
        index = index.byte()        
        index = index

        self.lamb = max(self.LambdaMin,self.LambdaMax/(1+0.1*self.it ))
        output = cos_theta * 1.0         
        output[index] -= cos_theta[index]*(1.0+0)/(1+self.lamb)        
        output[index] += phi_theta[index]*(1.0+0)/(1+self.lamb)        

        logpt = log_softmax(output)          
        logpt = ops.GatherD()(logpt, 1, target)
        logpt = logpt(target, (-1)) 
        pt = logpt.data.exp()

        loss = -1 * (1-pt)**self.gamma * logpt
        loss = loss.mean()   

        return loss
          

def get_loss(args):
    return AngleLoss(args)